#define _GNU_SOURCE 1
#include <iostream>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <cstring>
#include <cassert>
#include <cstdlib>
#include <poll.h>
#include <cerrno>
using namespace std;

const int USER_LIMIT = 5;   // 用户限制
const int BUFFER_SIZE = 64; // 缓冲区大小
const int FD_LIMIT = 65535; // 文件描述符限制

struct clientData {
    sockaddr_in address;    // 客户端地址
    char buf[BUFFER_SIZE];  // 读取缓冲区
    char* writeBuf;         // 写入缓冲区
};

int setNonBlocking(int fd) {
    int oldOption = fcntl(fd, F_GETFL);
    if (oldOption < 0) {
        return -1;
    }
    int newOption = oldOption | O_NONBLOCK;
    return fcntl(fd, F_SETFL, newOption);
}

int main(int argc, char* argv[])
{
    if (argc <= 2) {
        cerr << "Usage: " << argv[0] << " <IP> <Port>" << endl;
        return 1;
    }
    const char* ip = argv[1];
    int port = atoi(argv[2]);

    int ret = 0;
    struct sockaddr_in serverAddress;
    memset(&serverAddress, 0, sizeof(serverAddress));
    serverAddress.sin_family = AF_INET;
    inet_pton(AF_INET, ip, &serverAddress.sin_addr);
    serverAddress.sin_port = htons(port);

    int listenFd = socket(AF_INET, SOCK_STREAM, 0);
    assert(listenFd >= 0);

    ret = bind(listenFd, (struct sockaddr*)&serverAddress, sizeof(serverAddress));
    assert(ret != -1);

    ret = listen(listenFd, USER_LIMIT);
    assert(ret != -1);

    clientData* users = new clientData[FD_LIMIT]; // 用户数据
    pollfd fds[USER_LIMIT + 1]; // poll文件描述符数组
    int user_count = 0;         // 当前用户数量
    for(int i = 1; i <= USER_LIMIT; ++i) {
        fds[i].fd = -1;     // 初始化为无效文件描述符
        fds[i].events = 0;  // 初始化事件
    }
    fds[0].fd = listenFd;               // 监听套接字
    fds[0].events = POLLIN | POLLERR;   // 监听输入和错误事件
    fds[0].revents = 0;                 // 初始化返回事件

    while(1)
    {
        ret = poll(fds, user_count + 1, -1); // 阻塞等待事件
        if (ret < 0) {
            cerr << "poll failure: " << strerror(errno) << endl;
            break;
        }

        for (int i = 0; i < user_count + 1; ++i) {
            // 最开始时就只有一个listenFd，这个时候能读说明有新连接到来
            if (fds[i].fd == listenFd && (fds[i].revents & POLLIN)) {
                sockaddr_in clientAddr;
                socklen_t clientAddrLen = sizeof(clientAddr);
                int connFd = accept(listenFd, (struct sockaddr*)&clientAddr, &clientAddrLen);
                if (connFd < 0) {
                    cerr << "Accept failed: " << strerror(errno) << endl;
                    continue;
                }
                // 请求太多，则关闭新连接
                if (user_count >= USER_LIMIT) {
                    const char* info = "too many users\n";
                    cout << info << endl;
                    send(connFd, info, strlen(info), 0);
                    close(connFd);
                    continue;
                }

                user_count++;                       // 增加用户计数
                users[connFd].address = clientAddr; // 记录客户端地址
                setNonBlocking(connFd);             // 设置非阻塞模式
                fds[user_count].fd = connFd;        // 添加新连接到poll数组
                fds[user_count].events = POLLIN | POLLRDHUP | POLLERR; // 监听输入、关闭和错误事件
                fds[user_count].revents = 0;    // 初始化返回事件
                cout << "comes a new user, now we have " 
                     << user_count << " users" << endl;
            }
            else if (fds[i].revents & POLLERR) {
                cout << "get an error from " << fds[i].fd << endl;
                char errors[100];
                memset(errors, 0, sizeof(errors));
                socklen_t len = sizeof(errors);
                // 获取套接字错误信息
                if (getsockopt(fds[i].fd, SOL_SOCKET, SO_ERROR, &errors, &len) < 0) {
                    cerr << "get socket option failed: " << strerror(errno) << endl;
                }
                continue;
            }
            else if (fds[i].revents & POLLRDHUP) {
                // 处理客户端断开连接
                users[fds[i].fd] = users[fds[user_count].fd];
                close(fds[i].fd);
                fds[i] = fds[user_count];
                user_count--;
                i--;
                cout << "a client left, now we have " 
                     << user_count << " users" << endl;
                } 
            else if (fds[i].revents & POLLIN) {
                // 处理客户端发送的数据
                int connFd = fds[i].fd;
                memset(users[connFd].buf, 0, BUFFER_SIZE);
                ret = recv(connFd, users[connFd].buf, BUFFER_SIZE - 1, 0);
                if (ret < 0) {
                    if (errno != EAGAIN) {
                        close(connFd);
                        // 将当前断开连接的客户端对应的用户数据替换为最后一个用户的数据
                        users[fds[i].fd] = users[fds[user_count].fd];   
                        // 将pollfd数组中最后一个有效的文件描述符移动到当前断开连接的位置
                        fds[i] = fds[user_count];
                        user_count--;
                        // 下一次循环时重新检查当前索引位置
                        i--;
                    }
                }
                else if (ret == 0) {
                    // 对方关闭了连接，这个处理在EPOLLRDHUP处理了，因此这个不需要任何操作
                }
                else {
                    for(int j = 1; j <=user_count; ++j) {
                        if (fds[j].fd == connFd) {
                            continue; // 不发送给自己
                        }
                        fds[j].events |= POLLOUT;   // 设置可写事件
                        fds[j].events |= ~POLLIN;    // 清除可读事件
                        users[fds[j].fd].writeBuf = users[connFd].buf; // 将读取到的数据放入写入缓冲区
                    }
                }
            }
            else if (fds[i].revents & POLLOUT) {
                // 处理可写事件
                int connFd = fds[i].fd;
                if (!users[connFd].writeBuf) {
                    continue; // 没有数据可写
                }
                ret = send(connFd, users[connFd].writeBuf, strlen(users[connFd].writeBuf), 0);
                users[connFd].writeBuf = nullptr; // 清空写入缓冲区   
                
                fds[i].events |= ~POLLOUT; // 清除可写事件
                fds[i].events |= POLLIN;   // 恢复可读事件
            }
        }
    }
    delete[] users; // 释放用户数据
    close(listenFd);
    return 0;
}